# -*- coding: utf-8 -*-
"""
Dispatcher for RPC Objects. Routes messages and executes services.
"""
import asyncio
import logging
import six

from aiobsonrpc.concurrent import new_promise
from aiobsonrpc.definitions import RpcErrors
from aiobsonrpc.exceptions import JsonRpcError

__license__ = 'http://mozilla.org/MPL/2.0/'


class RpcForServices(object):
    """
    封装rpc对象, 限制外部方法调用
    """

    def __init__(self, rpc):
        self._rpc = rpc
        self._close_after = False
        self._aborted = False

    @property
    def aborted(self):
        return self._aborted

    @property
    def close_after_response_requested(self):
        return self._close_after

    def __getattr__(self, name):
        """
        禁止外部使用的方法
        """
        if name in ('close', 'join',) or name.startswith('_'):
            raise AttributeError(f"'{name}' is not allowed within service handler.'")
        return getattr(self._rpc, name)

    async def abort(self):
        await self._rpc.close()
        self._aborted = True

    def close_after_response(self):
        self._close_after = True


class Dispatcher(object):

    def __init__(self, rpc):
        """
        :param rpc: Rpc parent object.
        :type rpc: RpcBase
        """
        self._responses = {}  # {"<msg_id>": <promise>, ...}
        self._batch_responses = {}  # { ("<msg_id>", "<msg_id>",): promise, ...}
        self._active_task = []
        self.rpc = rpc
        self.conn_label = six.text_type(self.rpc.connection_id and '%s: ' % self.rpc.connection_id)
        self._task = self.rpc.loop.create_task(self.run())

    def __getattr__(self, name):
        return getattr(self.rpc, name)

    def _log_info(self, msg, *args, **kwargs):
        logging.info(self.conn_label + six.text_type(msg), *args, **kwargs)

    def _log_error(self, msg, *args, **kwargs):
        logging.error(self.conn_label + six.text_type(msg), *args, **kwargs)

    def register(self, msg_id):
        promise = new_promise()

        if isinstance(msg_id, tuple):
            self._batch_responses[msg_id] = promise
        else:
            self._responses[msg_id] = promise
        return promise

    def unregister(self, msg_id):
        if isinstance(msg_id, tuple):
            promise = self._batch_responses.get(msg_id)
            if msg_id in self._batch_responses:
                del self._batch_responses[msg_id]
        else:
            promise = self._responses.get(msg_id)
            if msg_id in self._responses:
                del self._responses[msg_id]
        if promise and not promise.is_set():
            promise.set(JsonRpcError('Timeout'))

    def _handle_parse_error(self, exception):
        try:
            self.rpc.socket_queue.put(
                self.rpc.definitions.error_response(
                    None, RpcErrors.parse_error, six.text_type(exception)))
        except:
            pass  # Effort made, success not required.
        self._log_error(exception)

    def _get_params(self, msg):
        """
        获取参数
        :param msg: dict
        """
        if 'params' not in msg:
            return [], {}
        params = msg['params']
        if isinstance(params, list):
            return params, {}
        if isinstance(params, dict):
            return [], params

    async def _execute_request(self, msg, rfs):
        """
        执行来自客户端的远程调用
        :param msg: dict
        """
        msg_id = msg['id']
        method_name = msg['method']
        args, kwargs = self._get_params(msg)
        try:
            method = self.rpc.services._request_handlers.get(method_name)  # 获取方法名
            if method:
                result = await method(self.rpc.services, rfs, *args, **kwargs)  # 类调用该方法
                return self.rpc.definitions.ok_response(msg_id, result)
            else:
                return self.rpc.definitions.error_response(
                    msg_id, RpcErrors.method_not_found)
        except Exception as e:
            return self.rpc.definitions.error_response(
                msg_id, RpcErrors.server_error, six.text_type(e))

    async def _handle_request(self, msg):
        """
        执行远程调用，回复响应写入数据流 response -> :net:
        :param msg: dict
        """
        self._log_info(u'Received request: ' + six.text_type(msg))
        rfs = RpcForServices(self.rpc)
        response = await self._execute_request(msg, rfs)
        if rfs.aborted:
            self._log_info(u'Connection aborted in request handler.')
            return
        await self.rpc.socket_queue.put(response)  # encode -> :net:
        self._log_info(u'Sent response: ' + six.text_type(response))
        if rfs.close_after_response_requested:
            await self.rpc.close()
            self._log_info(
                u'RPC closed due to invocation by Request handler.')

    async def _handle_batch_request(self, msg, rfs):
        """
        执行远程调用, 将 response 绑定给 event
        """
        self._log_info(u'Handling batch request: ' + six.text_type(msg))
        response = await self._execute_request(msg, rfs)
        self._log_info(
            u'Generated batch response: ' + six.text_type(response))
        promise = new_promise()
        promise.set(response)
        return promise.value

    async def _execute_notification(self, msg, rfs, after_effects):
        """
        执行来自客户端的通告方法
        :param msg: dict
        """
        method_name = msg['method']
        args, kwargs = self._get_params(msg)
        method = self.rpc.services._notification_handlers.get(method_name)
        if method:
            try:
                self._log_info(
                    u'Received notification: ' + six.text_type(msg))
                await method(self.rpc.services, rfs, *args, **kwargs)  # 执行通告方法
            except Exception as e:
                self._log_error(e)
            if (after_effects and not rfs.aborted and
                    rfs.close_after_response_requested):
                await self.rpc.close()
                self._log_info(
                    u'RPC closed due to invocation by '
                    u'Notification handler.')
        else:
            self._log_error(
                u'Unrecognized notification from peer: ' +
                six.text_type(msg))

    async def _handle_notification(self, msg):
        """
        调用通告方法处理程序
        :param msg: dict
        """
        rfs = RpcForServices(self.rpc)
        await self._execute_notification(msg, rfs, True)

    async def _handle_response(self, msg):
        """
        接收响应并绑定给对应的event (event.set())
        """
        msg_id = msg['id']
        promise = self._responses.get(msg_id)
        if promise:  # {"<msg_id>": <promise>, ...}
            if 'result' in msg:
                promise.set(msg['result'])  # 给事件绑定结果
            else:
                promise.set(RpcErrors.error_to_exception(msg['error']))
        else:
            self._log_error(
                u'Unrecognized/expired response from peer: ' +
                six.text_type(msg))

    async def _handle_nil_id_error_response(self, msg):
        """
        错误响应
        """
        self._log_error(msg)

    async def _handle_schema_error(self, msg):
        """
        dispatch 匹配不到的错误请求 error_response -> :net:
        """
        msg_id = None
        if isinstance(msg.get('id'), (six.string_types, int)):
            msg_id = msg['id']
        self.rpc.socket_queue.put(
            self.rpc.definitions.error_response(
                msg_id, RpcErrors.invalid_request))
        self._log_error(u'Invalid Request: ' + six.text_type(msg))

    async def _dispatch_batch(self, msgs):
        """
        遍历 msgs 处理请求|通告 response -> :net:
        """
        self._log_info(u'Received batch: ' + six.text_type(msgs))
        rfs = RpcForServices(self.rpc)
        promises = []
        notifies = []
        for msg in msgs:
            if 'id' in msg:
                promises.append(self._handle_batch_request(msg, rfs))  # 追加协程对象
            else:
                notifies.append(self._execute_notification(msg, rfs, False))
        results = await asyncio.gather(*promises)  # 等待协程对象执行完毕得到结果列表
        if results:
            if rfs.aborted:
                self._log_info(
                    'Connection aborted during batch processing.')
                return
            await self.rpc.socket_queue.put(results)  # encode -> :net:
            self._log_info(u'Sent batch response: ' + six.text_type(results))
        else:
            self._log_info(u'Notification-only batch processed.')
        if rfs.close_after_response_requested:
            await self.rpc.close()
            self._log_info(
                u'RPC closed due to invocation by Request or '
                u'Notification handler.')

    async def _handle_batch_response(self, msgs):
        """
        批量处理响应并绑定给对应的event (event.set())
        """

        def _extract_msg_content(msg):
            """
            返回响应数据 Ok or Error
            """
            if 'result' in msg:
                return msg['result']
            else:
                return RpcErrors.error_to_exception(msg['error'])

        without_id_msgs = list(filter(lambda x: x.get('id') is None, msgs))
        with_id_msgs = list(filter(lambda x: x.get('id') is not None, msgs))
        resp_map = dict(map(lambda x: (x['id'], x), with_id_msgs))  # {'id':msg, ...}
        msg_ids = set(resp_map.keys())
        resolved = False
        for id_tuple, promise in self._batch_responses.items():  # {(11, 12, 13, 14): promise}
            if msg_ids.issubset(set(id_tuple)):  # msg_ids <= set(id_tuple)
                batch_response = []
                for req_id in id_tuple:
                    if req_id in resp_map:
                        batch_response.append(
                            _extract_msg_content(resp_map[req_id]))
                    elif without_id_msgs:
                        batch_response.append(
                            _extract_msg_content(without_id_msgs.pop(0)))
                    else:
                        batch_response.append(
                            JsonRpcError(
                                'Peer did not respond to this request!'))
                promise.set(batch_response)
                resolved = True
                break
        if not resolved:
            self._log_error(
                u'Unrecognized/expired batch response from peer: ' +
                six.text_type(msgs))

    async def run(self):
        def _otherwise(*_):
            return True

        rpcd = self.rpc.definitions
        dispatch = {
            dict: [
                (rpcd.is_request, self._handle_request),
                (rpcd.is_notification, self._handle_notification),
                (rpcd.is_response, self._handle_response),
                (rpcd.is_nil_id_error_response,
                 self._handle_nil_id_error_response),
                (_otherwise, self._handle_schema_error),
            ],
            list: [
                (rpcd.is_batch_request, self._dispatch_batch),
                (rpcd.is_batch_response, self._handle_batch_response),
                (_otherwise, self._handle_schema_error),
            ]
        }

        self._log_info(u'---**---Start RPC message dispatcher---**---')
        while True:
            try:
                msg = await self.rpc.socket_queue.get()
                if msg is None:
                    break
                # self._log_info(u'Dispatch message.')
                if isinstance(msg, Exception):
                    self._handle_parse_error(msg)
                else:
                    for match_fn, handler_fn in dispatch.get(
                            type(msg),
                            [(_otherwise, self._handle_schema_error)]):  # （True,None)
                        if match_fn(msg):  # 判断 请求|通告|批量...
                            current_task = self.rpc.loop.create_task(handler_fn(msg))
                            self._active_task.append(current_task)
                            break
            except Exception as e:
                self._log_error(e)
        self._log_info(u'---**---Exit RPC message dispatcher---**---')

    async def join(self, timeout=None):
        def _is_alive(task):
            try:
                return not task.done()
            except AttributeError:
                return False

        self._active_task = list(filter(lambda t: _is_alive(t), self._active_task))
        if self._active_task:
            await asyncio.wait(self._active_task, timeout=timeout)
